import torch
import torch.nn as nn
from pointnet_util import farthest_point_sample, index_points, square_distance

# points: (B, npoint, channel) = [2, 1024, 64]
def sample_and_group(npoint, nsample, xyz, points):

    B, N, C = xyz.shape
    S = npoint 
    # 选出保留的 npoint 个点的 index
    fps_idx = farthest_point_sample(xyz, npoint) # [B, npoint]

    new_xyz = index_points(xyz, fps_idx) # [b, npoint, 3]
    new_points = index_points(points, fps_idx) # [b, npoint, 64], 得到 sample 的 point

    dists = square_distance(new_xyz, xyz)  # B x npoint x N, 计算 sample point 和所有 point 的距离
    idx = dists.argsort()[:, :, :nsample]  # B x npoint x nsample, 保留距离 sample point 最近的 K 个 point
    # 根据每个 group 的 index, 取出 每个 group 的 feature
    grouped_points = index_points(points, idx) # B * npoint * nsample *64
    grouped_points_norm = grouped_points - new_points.view(B, S, 1, -1) # 每个 sample point 的 K 近邻 feature 减去它自己
    # concat([2, 512, 32, 64], [2, 512, 32, 64], dim =-1) => [2, 512, 32, 128]
    # 每个 new point 的邻居的 feature，拼接 new point 本身的 feature，得到新的 feature
    new_points = torch.cat([grouped_points_norm, new_points.view(B, S, 1, -1).repeat(1, 1, nsample, 1)], dim=-1)
    # new_xyz: [B, npoint, 3]
    # new_points: [B, npoint, n_sample, 2*length]
    return new_xyz, new_points

# 聚合每个 group 的 feature
class Local_op(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size=1, bias=False)
        self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm1d(out_channels)
        self.bn2 = nn.BatchNorm1d(out_channels)
        self.relu = nn.ReLU()

    def forward(self, x):
        # new_feature: [B, npoint, n_sample, length] = [B, 512, 32, 128]
        b, n, s, d = x.size()  # torch.Size([B, 512, 32, 128])
        x = x.permute(0, 1, 3, 2) # [B, npoint, length, n_sample] # 因为最终是要聚合不同 sample 的 feature，因此这里把 length 维度作为 channel
        x = x.reshape(-1, d, s) # [B * npoint, length, n_sample]
        batch_size, _, N = x.size()
        x = self.relu(self.bn1(self.conv1(x))) # [B * npoint, length, n_sample] # 因为最终是要聚合不同 sample 的 feature，因此沿着不同 sample 的同一个 feature进行卷积操作
        x = self.relu(self.bn2(self.conv2(x))) # [B * npoint, length, n_sample]
        x = torch.max(x, 2)[0] # [B * npoint, length] # max pool 聚合操作，不同sample 的同一个 feature 取最大值
        x = x.view(batch_size, -1) # [B, length, npoint]
        x = x.reshape(b, n, -1).permute(0, 2, 1) # [B, length, npoint]
        return x


class SA_Layer(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.q_conv = nn.Conv1d(channels, channels // 4, 1, bias=False)
        self.k_conv = nn.Conv1d(channels, channels // 4, 1, bias=False)
        self.q_conv.weight = self.k_conv.weight 
        self.v_conv = nn.Conv1d(channels, channels, 1)
        self.trans_conv = nn.Conv1d(channels, channels, 1)
        self.after_norm = nn.BatchNorm1d(channels)
        self.act = nn.ReLU()
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x):
        # X: [B, length, npoint] = [B, 256, 256]
        x_q = self.q_conv(x).permute(0, 2, 1) # b, npoint, channel//4
        x_k = self.k_conv(x) # b, channel//4, npoint
        x_v = self.v_conv(x) # b, channel, npoint
        energy = x_q @ x_k # b, npoint, npoint
        attention = self.softmax(energy)
        attention = attention / (1e-9 + attention.sum(dim=1, keepdims=True))
        x_r = x_v @ attention # b, channel, npoint
        x_r = self.act(self.after_norm(self.trans_conv(x - x_r))) # this is offset attention
        x = x + x_r # skip connection
        return x
    

class StackedAttention(nn.Module):
    def __init__(self, channels=256):
        super().__init__()
        self.conv1 = nn.Conv1d(channels, channels, kernel_size=1, bias=False)
        self.conv2 = nn.Conv1d(channels, channels, kernel_size=1, bias=False)

        self.bn1 = nn.BatchNorm1d(channels)
        self.bn2 = nn.BatchNorm1d(channels)

        self.sa1 = SA_Layer(channels)
        self.sa2 = SA_Layer(channels)
        self.sa3 = SA_Layer(channels)
        self.sa4 = SA_Layer(channels)

        self.relu = nn.ReLU()
        
    def forward(self, x):
        # X: [B, length, npoint] = [B, 256, 256]
        # b, 3, npoint, nsample  
        # conv2d 3 -> 128 channels 1, 1
        # b * npoint, c, nsample 
        # permute reshape
        batch_size, _, N = x.size()

        x = self.relu(self.bn1(self.conv1(x))) # B, D, N
        x = self.relu(self.bn2(self.conv2(x)))

        x1 = self.sa1(x)
        x2 = self.sa2(x1)
        x3 = self.sa3(x2)
        x4 = self.sa4(x3)
        # x: [b, channel*4, npoint]
        x = torch.cat((x1, x2, x3, x4), dim=1)

        return x


class PointTransformerCls(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        output_channels = cfg.num_class
        self.conv1 = nn.Conv1d(cfg.input_dim, 64, kernel_size=1, bias=False)
        self.conv2 = nn.Conv1d(64, 64, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm1d(64)
        self.bn2 = nn.BatchNorm1d(64)
        self.gather_local_0 = Local_op(in_channels=128, out_channels=128)
        self.gather_local_1 = Local_op(in_channels=256, out_channels=256)
        self.pt_last = StackedAttention()

        self.relu = nn.ReLU()
        self.conv_fuse = nn.Sequential(nn.Conv1d(1280, 1024, kernel_size=1, bias=False),
                                   nn.BatchNorm1d(1024),
                                   nn.LeakyReLU(negative_slope=0.2))

        self.linear1 = nn.Linear(1024, 512, bias=False)
        self.bn6 = nn.BatchNorm1d(512)
        self.dp1 = nn.Dropout(p=0.5)
        self.linear2 = nn.Linear(512, 256)
        self.bn7 = nn.BatchNorm1d(256)
        self.dp2 = nn.Dropout(p=0.5)
        self.linear3 = nn.Linear(256, output_channels)

    def forward(self, x):
        # x: [2, 1024, 6]
        # xyz: [2, 1024, 3]
        xyz = x[..., :3]
        # x: [2, 6, 1024]  (B, input_channel, npoint)
        x = x.permute(0, 2, 1)
        batch_size, _, _ = x.size()
        x = self.relu(self.bn1(self.conv1(x))) # B, input_channel, length
        x = self.relu(self.bn2(self.conv2(x))) # B, input_channel, length = [B, channel, npoint]
        x = x.permute(0, 2, 1) # x: (B, npoint, channel) = [2, 1024, 64]
        # new_xyz: [B, npoint, 3]
        # new_feature: [B, npoint, n_sample, 2*length] = [B, 512, 32, 128]
        new_xyz, new_feature = sample_and_group(npoint=512, nsample=32, xyz=xyz, points=x)
        # 聚合每个 group 的 feature
        # feature_0: [B, 2*length, npoint]
        feature_0 = self.gather_local_0(new_feature)
        # feature: [B, npoint, 2*length]
        feature = feature_0.permute(0, 2, 1)
        # feature: [B, npoint, 2*length] = [B, 512, 32, 128]
        # new_feature: [B, npoint, 2*length] = [B, 256, 32, 256]
        new_xyz, new_feature = sample_and_group(npoint=256, nsample=32, xyz=new_xyz, points=feature)
        # 聚合每个 group 的 feature
        # feature_1: [B, channel, npoint] = [B, 256, 256]
        feature_1 = self.gather_local_1(new_feature)
        # 进行多层 attention 计算，并把每层的 attention 结果 concatenate
        # x: [B, channel*4, npoint]
        x = self.pt_last(feature_1) # Self Attention
        # 把 self attention 的结果和原始的 feature 进行 concat
        x = torch.cat([x, feature_1], dim=1) # x: [B, channel*5, npoint]
        # 一维卷积操作，在不同 point 的同一个 feature 进行卷积
        x = self.conv_fuse(x) # x: [B, channel*4, npoint]
        # 再进行 max pool 聚合操作，聚合到一个 channel
        x = torch.max(x, 2)[0] # x: [B, channel*4]
        #
        x = x.view(batch_size, -1)
        # FC 层
        x = self.relu(self.bn6(self.linear1(x)))
        x = self.dp1(x)
        x = self.relu(self.bn7(self.linear2(x)))
        x = self.dp2(x)
        x = self.linear3(x)

        return x